package com.ocean.web.util.nccalculate;

import ucar.ma2.*;
import ucar.nc2.*;

import java.io.IOException;
import java.time.Instant;
import java.util.*;

public class CalculateWindSpeedAndDirGRAPES {

    // 新文件的U/V分量变量名（GRAPES GFS格式）
    private static final String U_VAR_NAME = "UGRD_10maboveground";
    private static final String V_VAR_NAME = "VGRD_10maboveground";
    // 填充值（从新文件元数据中获取）
    private static final float FILL_VALUE = 9.999E20f;
    // 填充值容差（浮点数比较不能直接用==）
    private static final float FILL_VALUE_TOLERANCE = 1e15f;

    public static void main(String[] args) {
        // 输入文件路径（GRAPES GFS文件）
        String inputNcPath = "E:\\nc\\GRAPES\\CMA_BEN_GRAPESGFS_GLDT_2025111400_240h_SSW_cropped.nc";
        // 输出文件路径（仅含风速、风向的新文件） q
        String outputNcPath = "E:\\nc\\GRAPES\\CMA_BEN_GRAPESGFS_GLDT_2025111400_240h_SSW_calculate.nc";

        boolean success = CalculateWindSpeedAndDir(inputNcPath, outputNcPath);
        System.out.println(success ? "处理成功！新文件仅包含风速（wspd10）和风向（wdir10）" : "处理失败！");
    }

    /**
     * 仅计算10米高度风速（wspd10）和风向（wdir10），并生成独立文件
     */
    public static boolean CalculateWindSpeedAndDir(String inputNcPath, String outputNcPath) {
        try (NetcdfFile inputNc = NetcdfFile.open(inputNcPath)) {
            System.out.println("成功打开输入文件：" + inputNcPath);

            // 1. 读取U/V分量变量
            Variable uVar = inputNc.findVariable(U_VAR_NAME);
            Variable vVar = inputNc.findVariable(V_VAR_NAME);
            if (uVar == null || vVar == null) {
                throw new RuntimeException("输入文件中未找到U/V分量变量！U变量名：" + U_VAR_NAME + "，V变量名：" + V_VAR_NAME);
            }

            // 2. 校验维度一致性
            if (!uVar.getDimensions().equals(vVar.getDimensions())) {
                throw new RuntimeException("U/V分量变量维度不一致！");
            }

            // 3. 读取U/V数据和维度信息
            Array uArray = uVar.read();
            Array vArray = vVar.read();
            float[] uData = (float[]) uArray.copyTo1DJavaArray();
            float[] vData = (float[]) vArray.copyTo1DJavaArray();
            List<Dimension> dataDims = uVar.getDimensions(); // 维度：time[81] × latitude[224] × longitude[217]
            int[] dataShape = uVar.getShape(); // 维度形状数组

            // 4. 计算风速和风向（过滤填充值）
            int dataSize = uData.length;
            float[] wspd10Data = new float[dataSize];
            float[] wdir10Data = new float[dataSize];

            for (int i = 0; i < dataSize; i++) {
                float u = uData[i];
                float v = vData[i];

                // 无效数据处理
                if (isFillValue(u) || isFillValue(v)) {
                    wspd10Data[i] = FILL_VALUE;
                    wdir10Data[i] = FILL_VALUE;
                    continue;
                }

                // 风速：sqrt(u² + v²)
                wspd10Data[i] = (float) Math.sqrt(u * u + v * v);
                // 风向：风的来向（0°=北风，顺时针递增）
                double rad = Math.atan2(u, v);
                double deg = rad * 180 / Math.PI;
                wdir10Data[i] = (float) ((deg + 360) % 360);
            }

            // 5. 创建仅含风速、风向的输出文件
            createWindOnlyNcFile(inputNc, dataDims, wspd10Data, wdir10Data, dataShape, outputNcPath);
            System.out.println("计算完成！新文件路径：" + outputNcPath);
            return true;

        } catch (Exception e) {
            System.err.println("处理失败：" + e.getMessage());
            e.printStackTrace();
            return false;
        }
    }

    /**
     * 判断是否为填充值
     */
    private static boolean isFillValue(float value) {
        return Math.abs(value - FILL_VALUE) < FILL_VALUE_TOLERANCE;
    }

    /**
     * 创建仅含风速、风向的NetCDF文件（删除所有其他变量）
     */
    private static void createWindOnlyNcFile(NetcdfFile inputNc, List<Dimension> dataDims,
                                             float[] wspd10Data, float[] wdir10Data, int[] dataShape, String outputPath) throws IOException {

        NetcdfFileWriter writer = NetcdfFileWriter.createNew(NetcdfFileWriter.Version.netcdf4, outputPath);
        Group rootGroup = null; // 根组（低版本兼容）

        try {
            // 1. 复制核心维度（time、latitude、longitude）
            Map<String, Dimension> createdDims = new HashMap<>();
            for (Dimension dim : dataDims) {
                String dimName = dim.getShortName();
                // 从原文件获取维度的完整属性（如latitude的units、axis等）
                Variable dimVar = inputNc.findVariable(dimName); // 维度变量（如latitude变量本身）
                Dimension newDim = writer.addDimension(rootGroup, dimName, dim.getLength());
                createdDims.put(dimName, newDim);

                // 复制维度变量的属性（如latitude的degrees_north、axis=Y等）
                if (dimVar != null) {
                    Variable newDimVar = writer.addVariable(rootGroup, dimName, dimVar.getDataType(), Collections.singletonList(newDim));
                    copyVariableAttributes(dimVar, newDimVar);
                }
            }

            // 2. 创建风速变量（wspd10）
            Variable wspd10Var = writer.addVariable(rootGroup, "wspd10", DataType.FLOAT, dataDims);
            //wspd10Var.addAttribute(new Attribute("rwid", "RW_"+Instant.now().toEpochMilli()));
            wspd10Var.addAttribute(new Attribute("long_name", "10米高度风速"));
            wspd10Var.addAttribute(new Attribute("units", "m/s"));
            wspd10Var.addAttribute(new Attribute("_FillValue", FILL_VALUE));
            wspd10Var.addAttribute(new Attribute("missing_value", FILL_VALUE));
            wspd10Var.addAttribute(new Attribute("short_name", "wspd10"));
            wspd10Var.addAttribute(new Attribute("level", "10 m above ground"));

            // 3. 创建风向变量（wdir10）
            Variable wdir10Var = writer.addVariable(rootGroup, "wdir10", DataType.FLOAT, dataDims);
            //wdir10Var.addAttribute(new Attribute("rwid", "RW_"+Instant.now().toEpochMilli()));
            wdir10Var.addAttribute(new Attribute("long_name", "10米高度风向（风的来向）"));
            wdir10Var.addAttribute(new Attribute("units", "degree"));
            wdir10Var.addAttribute(new Attribute("_FillValue", FILL_VALUE));
            wdir10Var.addAttribute(new Attribute("missing_value", FILL_VALUE));
            wdir10Var.addAttribute(new Attribute("short_name", "wdir10"));
            wdir10Var.addAttribute(new Attribute("level", "10 m above ground"));

            // 4. 复制必要的全局属性（确保文件可解析）
            List<String> necessaryGlobalAttrs = Arrays.asList("Conventions", "history", "History", "CDI", "CDO");
            for (Attribute attr : inputNc.getGlobalAttributes()) {
                if (necessaryGlobalAttrs.contains(attr.getShortName())) {
                    writer.addGlobalAttribute(attr);
                }
            }
            // 添加自定义全局属性，说明文件来源
            writer.addGlobalAttribute(new Attribute("source_file", inputNc.getLocation()));
            writer.addGlobalAttribute(new Attribute("created_by", "Java NetCDF 5.4.1"));
            writer.addGlobalAttribute(new Attribute("variable_description", "仅包含10米风速（wspd10）和风向（wdir10）"));

            // 5. 写入数据并关闭文件
            writer.create();
            // 写入维度数据（如latitude、longitude、time的具体值）
            for (Dimension dim : dataDims) {
                String dimName = dim.getShortName();
                Variable srcDimVar = inputNc.findVariable(dimName);
                Variable destDimVar = writer.findVariable(dimName);
                if (srcDimVar != null && destDimVar != null) {
                    writer.write(destDimVar, srcDimVar.read());
                }
            }
            // 写入风速、风向数据
            writer.write(wspd10Var, Array.factory(DataType.FLOAT, dataShape, wspd10Data));
            writer.write(wdir10Var, Array.factory(DataType.FLOAT, dataShape, wdir10Data));

            System.out.println("文件创建成功！新文件仅包含以下变量：");
            System.out.println("- 维度变量：time, latitude, longitude");
            System.out.println("- 数据变量：wspd10（风速）, wdir10（风向）");

        } catch (InvalidRangeException e) {
            throw new RuntimeException(e);
        } finally {
            writer.close();
        }
    }

    /**
     * 复制变量属性
     */
    private static void copyVariableAttributes(Variable srcVar, Variable destVar) {
        for (Attribute attr : srcVar.getAttributes()) {
            destVar.addAttribute(attr);
        }
    }
}